"""
mathmodeling_SAK.Analysis.ARIMA
===============================

定义了 ARIMA 模型需要用到的工具。
"""

import pandas as pd
import matplotlib.pyplot as plt

import seaborn as sns

import statsmodels.api as sm
from statsmodels.tsa.arima.model import ARIMA

from statsmodels.tsa.stattools import acf, pacf
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf

from .trend_testing import autoADF
from .trend_testing import iterative_trend_tests

from .evalution import evaluate_model_df

class ARIMA_AllInOne:
    """
    封装好的 ARIMA 模型类，基于 `statsmodels.tsa.arima`
    一键进行便捷的 ARIMA 模型的数据趋势性检验、定阶、拟合、可视化以及
    模型拟合优度检验，以便进行结果分析。

    输入参数
    ----------
    data : pd.DataFrame
        包含时间序列数据的数据框。
    time_column_name : str
        时间列的列名。
    value_column_name : str
        值列的列名。
    alpha : float, optional
        显著性水平，该数值主要用于数据平稳性的 adf 检验中与 p-values
        进行比较，默认值为0.05。
    nlags : int, optional
        ACF和PACF图的滞后阶数，默认值为10。
    trend_test_method : str, optional
        趋势检验方法，默认值为'Mann-Kendall'。
    max_iter_of_trend_test : int, optional
        差分次数检验的最大迭代次数，可选参数有 'All', 'Cox-Stuart', 
        'Mann-Kendall', 'Slope'。如果选择了 'All' 则要求所有检验
        都认为数据无趋势时，迭代才会停止。 默认为 'Mann-Kendall'。
        需要注意的是，在这个类的定义当中，对数据趋势性的检验依赖的是单
        位根检验（ADF）方法。这个参数仅仅在判断模型差分次数的时候才会起
        作用。
        
    使用示例
    ----------
    
    读取数据后传入 Data
    
    >>> arima = ARIMA_AllInOne(
            data,
            time_column_name='日期',
            value_column_name='花叶类'
        )
        
    调用 `.do_everything` 方法之后，这个类将会帮你解决其余全部的问题。
    
    >>> arima.do_everything()
    """
    def __init__(
        self, data, 
        time_column_name, 
        value_column_name,
        alpha=0.05,
        nlags=10,
        trend_test_method='Mann-Kendall',
        max_iter_of_trend_test=5
        ):
        self.data = data
        self.t = time_column_name
        self.v = value_column_name
        self.trend_test_method = trend_test_method
        self.trend_test_max_iter = max_iter_of_trend_test
        self.diff_level = 0
        self.trend = True
        self.alpha = alpha
        self.nlags = nlags
        
    def adf(self):
        """ 进行ADF检验以确定数据的平稳性，并根据检验结果设置趋势属性。
        """
        self.adf_result = autoADF(self.data[self.v])
        self.adf_p_value = self.adf_result.loc[
            'p-value', 
            'value']
        if self.adf_p_value <= self.alpha:
            self.trend = False
            
    def get_diff_n(self):
        """ 确定数据平稳化所需的差分次数。
        """
        if self.trend == False:
            self.diff_level = 0
        else:
            self.diff_level = iterative_trend_tests(
                self.data[self.v],
                trend_test_method=self.trend_test_method,
                max_iter_of_trend_test=self.trend_test_max_iter
                )['Iteration'].max()
        
    def acf_pacf(self, nlags=10):
        """
        计算并存储自相关函数（ACF）和偏自相关函数（PACF）。

        输入参数
        ----------
        nlags : int, optional
            ACF和PACF的滞后阶数，默认值为10。
        """
        self.lag_acf = acf(
            self.data[self.v], nlags=nlags)
        self.lag_pacf = pacf(
            self.data[self.v], nlags=nlags, method='ols')
        
    def plot_acf_pacf(
        self, nlags=10, 
        fig_kwgs={}
        ):
        """
        绘制ACF和PACF图。

        输入参数
        ----------
        nlags : int, optional
            ACF和PACF的滞后阶数，默认值为10。
        fig_kwgs : dict, optional
            用于plt.subplots的参数，默认值为空字典。

        返回值
        ----------
        fig : matplotlib.figure.Figure
            绘制的图形对象。
        """
        fig, axes = plt.subplots(2, 1, **fig_kwgs)
        plot_acf(self.data[self.v], lags=nlags, ax=axes[0])
        plot_pacf(self.data[self.v], lags=nlags, ax=axes[1])
        self.acf_pacf_plot = fig
        plt.tight_layout()
        # return fig
    
    def ic(self, max_ar=3, max_ma=3):
        """
        计算不同 ARIMA 模型的 AIC、BIC 和 HQIC 值以选择最佳模型。

        输入参数
        ----------
        max_ar : int, optional
            最大自回归阶数，默认值为3。
        max_ma : int, optional
            最大移动平均阶数，默认值为3。

        返回值
        -------
        无
        """
        if self.trend == True:
            trend = 'n'
        else:
            trend = 'c'
        self.ic_result = sm.tsa.arma_order_select_ic(
            self.data[self.v], 
            max_ar = max_ar, 
            max_ma = max_ma,
            ic = ['aic', 'bic', 'hqic' ],
            trend = trend, 
            model_kw = None, 
            fit_kw = None
        )
        
    def plot_ics_heatmap(
        self,
        fig_kwgs={
            'figsize' : (17, 4)
        }
        ):
        """
        绘制AIC、BIC和HQIC的热力图。

        输入参数
        ----------
        figure_kwargs : dict, optional
            用于plt.subplots的参数，默认值包含'figsize'。

        返回值
        ----------
        fig : matplotlib.figure.Figure
            绘制的图形对象。
        """
        fig, axes = plt.subplots(1, 3, **fig_kwgs)
        sns.heatmap(self.ic_result['aic'], ax=axes[0])
        sns.heatmap(self.ic_result['bic'], ax=axes[1])
        sns.heatmap(self.ic_result['hqic'], ax=axes[2])
        axes[0].set_title('AIC')
        axes[1].set_title('BIC')
        axes[2].set_title('HQIC')
        self.ics_heatmaps = fig
        # return fig
    
    def fit_model(self, order=(1, 1, 1)):
        """
        拟合ARIMA模型。

        输入参数
        ----------
        order : tuple, optional
            ARIMA模型的(p, d, q)阶数，默认值为(1, 1, 1)。

        返回值
        ----------
        result : statsmodels.tsa.arima.model.ARIMAResults
            拟合结果。
        model : statsmodels.tsa.arima.model.ARIMA
            拟合的模型对象。
        """
        model = ARIMA(
            self.data[self.v], 
            order=order)
        result = model.fit()
        return result, model
    
    def fit_all3_models(self):
        """ 拟合基于AIC、BIC和HQIC的三个最佳ARIMA模型，并存储结果。
        """
        self.result_dict_list = []
        
        ics_name_list = [
            'aic_min_order',
            'bic_min_order',
            'hqic_min_order'
        ]
        
        for ic in ics_name_list:
            
            result_dict = {
                'ic' : None,
                'order' : None,
                'result' : None,
                'model' : None,
                'evaluations' : None
            }
            
            order = (
                self.ic_result[ic][0],
                self.diff_level,
                self.ic_result[ic][1]
            )
            result, model = self.fit_model(
                order=order
            )
            
            result_dict['ic'] = ic
            result_dict['order'] = order
            result_dict['result'] = result
            result_dict['model'] = model
            
            result_dict['evaluations'] = evaluate_model_df(
                self.data[self.v], result.fittedvalues)
            
            self.result_dict_list.append(result_dict)
            
    def plot_fit(
        self, 
        fig_kwgs={
            'figsize' : (20, 4)
            },
        line_kwgs = {}
        ):
        """
        绘制原始数据和拟合值的对比图。

        输入参数
        ----------
        fig_kwgs : dict, optional
            用于plt.subplots的参数，默认值包含'figsize'。
        line_kwgs : dict, optional
            用于sns.lineplot的参数，默认值为空字典。

        返回值
        ----------
        fig : matplotlib.figure.Figure
            绘制的图形对象。
        """
        fig, axes = plt.subplots(1, 3, **fig_kwgs)
        for i in range(3):
            sns.lineplot(
                x=self.data[self.t],
                y=self.data[self.v], ax=axes[i],
                label='Original Data', **line_kwgs
            )
            sns.lineplot(
                x=self.data[self.t],
                y=self.result_dict_list[i]['result'
                    ].fittedvalues,
                ax=axes[i], label='Fitted Values',
                **line_kwgs
            )
            axes[i].set_xticks(self.data[self.t])
            axes[i].set_xticklabels(
                labels=self.data[self.t], rotation=90)
            axes[i].set_title(
                self.result_dict_list[i]['ic'])
            axes[i].legend()
        
        self.fit_plot = fig
        # return fig
    
    def combine_evaluations(self):
        """ 合并所有模型的评估结果。
        """
        combined_df = pd.concat(
            [
                result_dict[
                    'evaluations'
                    ] for result_dict in self.result_dict_list
                ], axis=1
            )
        combined_df.columns = [
            'aic', 'bic', 'hqic'
        ]
        self.evaluations_df = combined_df
    
    def plot_metrics(
        self,
        fig_kwgs={
            'figsize' : (20, 8)
            },
        bar_kwgs={
            'palette' : 'Paired'
            }
        ):
        """ 绘制不同准则下各项模型拟合优度评价指标的柱状对比图
        """
        fig, axes = plt.subplots(2, 3, **fig_kwgs)
        evaluations_df_colnames = self.evaluations_df.T.columns
        for i in range(1, 7):
            sns.barplot(
                data=self.evaluations_df.T.reset_index()[[
                    'index', evaluations_df_colnames[i - 1]
                    ]].melt(
                        id_vars='index', 
                        var_name='varible', 
                        value_name='value'),
                x='varible', y='value', hue='index',
                ax=axes[(i-1) // 3][(i-1) % 3], **bar_kwgs
            )
        plt.tight_layout()
        self.metrics_plot = fig
        # return fig
    
    def do_everything(self):
        """ 一键完成所有任务
        """
        self.adf()
        self.ic()
        self.fit_all3_models()
        
        # self.plot_acf_pacf()
        # self.plot_ics_heatmap()
        # self.plot_fit()
        
        self.combine_evaluations()
        
        print('models are fitted.')
        print('the following methods can be used:')
        
        print('    plot_acf_pacf()')
        print('    plot_ics_heatmap()')
        print('    plot_fit()')
        print('    plot_metrics()')
        
        # self.plot_metrics()
